{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Copyright 2020 The OpenFermion Developers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FQE vs OpenFermion vs Cirq: Diagonal Coulomb Operators"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://quantumai.google/openfermion/fqe/tutorials/diagonal_coulomb_evolution\"><img src=\"https://quantumai.google/site-assets/images/buttons/quantumai_logo_1x.png\" />View on QuantumAI</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/quantumlib/OpenFermion/blob/master/docs/fqe/tutorials/diagonal_coulomb_evolution.ipynb\"><img src=\"https://quantumai.google/site-assets/images/buttons/colab_logo_1x.png\" />Run in Google Colab</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://github.com/quantumlib/OpenFermion/blob/master/docs/fqe/tutorials/diagonal_coulomb_evolution.ipynb\"><img src=\"https://quantumai.google/site-assets/images/buttons/github_logo_1x.png\" />View source on GitHub</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://storage.googleapis.com/tensorflow_docs/OpenFermion/docs/fqe/tutorials/diagonal_coulomb_evolution.ipynb\"><img src=\"https://quantumai.google/site-assets/images/buttons/download_icon_1x.png\" />Download notebook</a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Special routines are available for evolving under a diagonal Coulomb operator.  This notebook describes how to use these built in routines and how they work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import fqe\n",
    "except ImportError:\n",
    "    !pip install fqe --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import product\n",
    "import fqe\n",
    "from fqe.hamiltonians.diagonal_coulomb import DiagonalCoulomb\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import openfermion as of\n",
    "\n",
    "from scipy.linalg import expm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Utility function\n",
    "def uncompress_tei(tei_mat, notation='chemistry'):\n",
    "    \"\"\"\n",
    "    uncompress chemist notation integrals\n",
    "\n",
    "    tei_tensor[i, k, j, l] = tei_mat[(i, j), (k, l)]\n",
    "    [1, 1, 2, 2] = [1, 1, 2, 2] = [1, 1, 2, 2]  = [1, 1, 2, 2]\n",
    "    [i, j, k, l] = [k, l, i, j] = [j, i, l, k]* = [l, k, j, i]*\n",
    "\n",
    "    For real we also have swap of i <> j and k <> l\n",
    "    [j, i, k, l] = [l, k, i, j] = [i, j, l, k] = [k, l, j, i]\n",
    "\n",
    "    tei_mat[(i, j), (k, l)] = int dr1 int dr2 phi_i(dr1) phi_j(dr1) O(r12) phi_k(dr1) phi_l(dr1)\n",
    "\n",
    "    Physics notation is the notation that is used in FQE.\n",
    "\n",
    "    Args:\n",
    "        tei_mat: compressed two electron integral matrix\n",
    "\n",
    "    Returns:\n",
    "        uncompressed 4-electron integral tensor. No antisymmetry.\n",
    "    \"\"\"\n",
    "    if notation not in ['chemistry', 'physics']:\n",
    "        return ValueError(\"notation can be [chemistry, physics]\")\n",
    "\n",
    "    norbs = int(0.5 * (np.sqrt(8 * tei_mat.shape[0] + 1) - 1))\n",
    "    basis = {}\n",
    "    cnt = 0\n",
    "    for i, j in product(range(norbs), repeat=2):\n",
    "        if i >= j:\n",
    "            basis[(i, j)] = cnt\n",
    "            cnt += 1\n",
    "\n",
    "    tei_tensor = np.zeros((norbs, norbs, norbs, norbs))\n",
    "    for i, j, k, l in product(range(norbs), repeat=4):\n",
    "        if i >= j and k >= l:\n",
    "            tei_tensor[i, j, k, l] = tei_mat[basis[(i, j)], basis[(k, l)]]\n",
    "            tei_tensor[k, l, i, j] = tei_mat[basis[(i, j)], basis[(k, l)]]\n",
    "            tei_tensor[j, i, l, k] = tei_mat[basis[(i, j)], basis[(k, l)]]\n",
    "            tei_tensor[l, k, j, i] = tei_mat[basis[(i, j)], basis[(k, l)]]\n",
    "\n",
    "            tei_tensor[j, i, k, l] = tei_mat[basis[(i, j)], basis[(k, l)]]\n",
    "            tei_tensor[l, k, i, j] = tei_mat[basis[(i, j)], basis[(k, l)]]\n",
    "            tei_tensor[i, j, l, k] = tei_mat[basis[(i, j)], basis[(k, l)]]\n",
    "            tei_tensor[k, l, j, i] = tei_mat[basis[(i, j)], basis[(k, l)]]\n",
    "\n",
    "    if notation == 'chemistry':\n",
    "        return tei_tensor\n",
    "    elif notation == 'physics':\n",
    "        return np.asarray(tei_tensor.transpose(0, 2, 1, 3), order='C')\n",
    "\n",
    "    return tei_tensor\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first example we will perform is diagonal Coulomb evolution on the Hartree-Fock state.  The diagonal Coulomb operator is defined as\n",
    "\n",
    "\\begin{align}\n",
    "V = \\sum_{\\alpha, \\beta \\in \\{\\uparrow, \\downarrow\\}}\\sum_{p,q} V_{pq,pq}n_{p,\\alpha}n_{q,\\beta}\n",
    "\\end{align}\n",
    "\n",
    "The number of free parpameters are $\\mathcal{O}(N^{2})$ where $N$ is the rank of the spatial basis. The `DiagonalCoulomb` Hamiltonian takes either a generic 4-index tensor or the $N \\times N$ matrix defining $V$.  If the 4-index tensor is given the $N \\times N$ matrix is constructed along with the diagonal correction.  If the goal is to just evolve under $V$ it is recommended the user input the $N \\times N$ matrix directly.\n",
    "\n",
    "All the terms in $V$ commute and thus we can evolve under $V$ exactly by counting the accumulated phase on each bitstring.\n",
    "\n",
    "\n",
    "To start out let's define a Hartree-Fock wavefunction for 4-orbitals and 2-electrons $S_{z} =0$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "norbs = 4\n",
    "tedim = norbs * (norbs + 1) // 2\n",
    "if (norbs // 2) % 2 == 0:\n",
    "    n_elec = norbs // 2\n",
    "else:\n",
    "    n_elec = (norbs // 2) + 1\n",
    "sz = 0\n",
    "fqe_wfn = fqe.Wavefunction([[n_elec, sz, norbs]])\n",
    "fci_data = fqe_wfn.sector((n_elec, sz))\n",
    "fci_graph = fci_data.get_fcigraph()\n",
    "hf_wf = np.zeros((fci_data.lena(), fci_data.lenb()), dtype=np.complex128)\n",
    "hf_wf[0, 0] = 1  # right most bit is zero orbital.\n",
    "fqe_wfn.set_wfn(strategy='from_data',\n",
    "                raw_data={(n_elec, sz): hf_wf})\n",
    "fqe_wfn.print_wfn()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can define a random 2-electron operator $V$.  To define $V$ we need a $4 \\times 4$ matrix.  We will generate this matrix by making a full random two-electron integral matrix and then just take the diagonal elements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tei_compressed = np.random.randn(tedim**2).reshape((tedim, tedim))\n",
    "tei_compressed = 0.5 * (tei_compressed + tei_compressed.T)\n",
    "tei_tensor = uncompress_tei(tei_compressed, notation='physics')\n",
    "\n",
    "diagonal_coulomb = of.FermionOperator()\n",
    "diagonal_coulomb_mat = np.zeros((norbs, norbs))\n",
    "for i, j in product(range(norbs), repeat=2):\n",
    "    diagonal_coulomb_mat[i, j] = tei_tensor[i, j, i, j]\n",
    "    for sigma, tau in product(range(2), repeat=2):\n",
    "        diagonal_coulomb += of.FermionOperator(\n",
    "            ((2 * i + sigma, 1), (2 * i + sigma, 0), (2 * j + tau, 1),\n",
    "             (2 * j + tau, 0)), coefficient=diagonal_coulomb_mat[i, j])\n",
    "\n",
    "dc_ham = DiagonalCoulomb(diagonal_coulomb_mat)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evolution under $V$ can be computed by looking at each bitstring, seeing if $n_{p\\alpha}n_{q\\beta}$ is non-zero and then phasing that string by $V_{pq}$.  For the Hartree-Fock state we can easily calculate this phase accumulation.  The alpha and beta bitstrings are \"0001\" and \"0001\". "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha_occs = [list(range(fci_graph.nalpha()))]\n",
    "beta_occs = [list(range(fci_graph.nbeta()))]\n",
    "occs = alpha_occs[0] + beta_occs[0]\n",
    "diag_ele = 0.\n",
    "for ind in occs:\n",
    "    for jnd in occs:\n",
    "        diag_ele += diagonal_coulomb_mat[ind, jnd]\n",
    "evolved_phase = np.exp(-1j * diag_ele)\n",
    "print(evolved_phase)\n",
    "\n",
    "# evolve FQE wavefunction\n",
    "evolved_hf_wfn = fqe_wfn.time_evolve(1, dc_ham)\n",
    "\n",
    "# check they the accumulated phase is equivalent!\n",
    "assert np.isclose(evolved_hf_wfn.get_coeff((n_elec, sz))[0, 0], evolved_phase)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now try this out for more than 2 electrons.  Let's reinitialize a wavefunction on 6-orbitals with 4-electrons $S_{z} = 0$ to a random state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "norbs = 6\n",
    "tedim = norbs * (norbs + 1) // 2\n",
    "if (norbs // 2) % 2 == 0:\n",
    "    n_elec = norbs // 2\n",
    "else:\n",
    "    n_elec = (norbs // 2) + 1\n",
    "sz = 0\n",
    "fqe_wfn = fqe.Wavefunction([[n_elec, sz, norbs]])\n",
    "fqe_wfn.set_wfn(strategy='random')\n",
    "initial_coeffs = fqe_wfn.get_coeff((n_elec, sz)).copy()\n",
    "print(\"Random initial wavefunction\")\n",
    "fqe_wfn.print_wfn()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to build our Diagoanl Coulomb operator For this bigger system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tei_compressed = np.random.randn(tedim**2).reshape((tedim, tedim))\n",
    "tei_compressed = 0.5 * (tei_compressed + tei_compressed.T)\n",
    "tei_tensor = uncompress_tei(tei_compressed, notation='physics')\n",
    "\n",
    "diagonal_coulomb = of.FermionOperator()\n",
    "diagonal_coulomb_mat = np.zeros((norbs, norbs))\n",
    "for i, j in product(range(norbs), repeat=2):\n",
    "    diagonal_coulomb_mat[i, j] = tei_tensor[i, j, i, j]\n",
    "    for sigma, tau in product(range(2), repeat=2):\n",
    "        diagonal_coulomb += of.FermionOperator(\n",
    "            ((2 * i + sigma, 1), (2 * i + sigma, 0), (2 * j + tau, 1),\n",
    "             (2 * j + tau, 0)), coefficient=diagonal_coulomb_mat[i, j])\n",
    "\n",
    "dc_ham = DiagonalCoulomb(diagonal_coulomb_mat)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can convert our wavefunction to a cirq wavefunction, evolve under the diagonal_coulomb operator we constructed and then compare the outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cirq_wfn = fqe.to_cirq(fqe_wfn).reshape((-1, 1))\n",
    "final_cirq_wfn = expm(-1j * of.get_sparse_operator(diagonal_coulomb).todense()) @ cirq_wfn\n",
    "# recover a fqe wavefunction\n",
    "from_cirq_wfn = fqe.from_cirq(final_cirq_wfn.flatten(), 1.0E-8)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fqe_wfn = fqe_wfn.time_evolve(1, dc_ham)\n",
    "print(\"Evolved wavefunction\")\n",
    "fqe_wfn.print_wfn()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"From Cirq Evolution\")\n",
    "from_cirq_wfn.print_wfn()\n",
    "assert np.allclose(from_cirq_wfn.get_coeff((n_elec, sz)),\n",
    "                   fqe_wfn.get_coeff((n_elec, sz)))\n",
    "print(\"Wavefunctions are equivalent\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can compare against evolving each term of $V$ individually."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fqe_wfn = fqe.Wavefunction([[n_elec, sz, norbs]])\n",
    "fqe_wfn.set_wfn(strategy='from_data',\n",
    "                raw_data={(n_elec, sz): initial_coeffs})\n",
    "for term, coeff in diagonal_coulomb.terms.items():\n",
    "    op = of.FermionOperator(term, coefficient=coeff)\n",
    "    fqe_wfn = fqe_wfn.time_evolve(1, op)\n",
    "\n",
    "assert np.allclose(from_cirq_wfn.get_coeff((n_elec, sz)),\n",
    "               fqe_wfn.get_coeff((n_elec, sz)))\n",
    "print(\"Individual term evolution is equivalent\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
